import os, sys
import time
import queue
from astropy.io import fits
import numpy as np
from scipy.stats import trim_mean, sigmaclip

from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QPushButton,
    QTextEdit, QLabel, QProgressBar, QGridLayout, QGroupBox
)
from PySide6.QtCore import (
    QThread, Signal, QMutex, QWaitCondition, Qt
)

from ccd.ccd_config import (
    ccdDataConfig, ccdDataFilter,
    ccdPreOverScanCorrection,
    ccdOverScanCorrection,
    estimateGainUsingCorrectedFlats
)

class WorkerThread_compareFlats(QThread):
    """
    This class is used to compare vertical bright streak in flats.
    """
    # 定义信号：用于传递字符串消息

    # 给主线程发送消息
    signal_message          = Signal(str)
    signal_result           = Signal(dict)
    signal_plot_fro_gain    = Signal(dict)  #
    signal_plot_flat_ratio  = Signal(dict)
    signal_finished         = Signal()

    # 从主线程接受参数利用__init__()，见下面：

    def __init__(
            self,
            thread_id,
            wkdir=None,
            flatFilename1=None,
            flatFilename2=None,
            biasFilename1=None,
            biasFilename2=None,
            applyGainCorrection=False,
            applySelfNorm=False,
            showCorrectedFlats=False,
            showCorrectedBias=False,
            anchorToZero=False,
            correctExpTime=False,
            expTime1=None,
            expTime2=None,
            plotFlatRatio=False,
            useColPreScan=False,
            parent=None
        ):
        super().__init__(parent)
        self.thread_id          = thread_id
        self.wkdir              = wkdir
        self.flatFilename1      = flatFilename1
        self.flatFilename2      = flatFilename2
        self.biasFilename1      = biasFilename1
        self.biasFilename2      = biasFilename2
        self.applyGainCorrection= applyGainCorrection
        self.applySelfNorm      = applySelfNorm
        self.showCorrectedFlats = showCorrectedFlats
        self.showCorrectedBias  = showCorrectedBias
        self.anchorToZero       = anchorToZero
        self.correctExpTime     = correctExpTime
        self.expTime1           = expTime1
        self.expTime2           = expTime2
        self.plotFlatRatio      = plotFlatRatio
        self.useColPreScan      = useColPreScan

        self.mutex              = QMutex()
        self.condition          = QWaitCondition()
        self.queue              = queue.Queue()

    def run(self):
        # self.signal_message.emit(f"start comparing flats, my thread_id is: {self.thread_id}")
        self.signal_message.emit(f"start comparing flats...")

        corr_flat1 = '{}/corrected_flat1.fits'.format(self.wkdir)
        corr_flat2 = '{}/corrected_flat2.fits'.format(self.wkdir)
        corr_bias1 = '{}/corrected_bias1.fits'.format(self.wkdir)
        corr_bias2 = '{}/corrected_bias2.fits'.format(self.wkdir)

        flat1 = None
        flat2 = None
        bias1 = None
        bias2 = None
        fro_gain1 = None
        fro_gain2 = None
        bro_gain1 = None
        bro_gain2 = None

        if self.biasFilename1 is not None:
            if os.path.isfile(self.biasFilename1):

                if self.useColPreScan:
                    bias1, bro_gain1 = ccdPreOverScanCorrection(self.biasFilename1, applyGainCorrection=self.applyGainCorrection)
                else:
                    bias1, bro_gain1 = ccdOverScanCorrection(self.biasFilename1, applyGainCorrection=self.applyGainCorrection)

                if self.showCorrectedBias:
                    hdu = fits.PrimaryHDU(bias1)
                    hdu.writeto(corr_bias1,overwrite=True)
                    self.signal_message.emit('saved {}'.format(corr_bias1))
                self.signal_message.emit('got row-over/col-pre scan corrected bias from: {}'.format(self.biasFilename1))
            else:
                self.signal_message.emit("bias1: '{}' does not exist".format(self.biasFilename1))

        if self.biasFilename2 is not None:
            if os.path.isfile(self.biasFilename2):

                if self.useColPreScan:
                    bias2, bro_gain2 = ccdPreOverScanCorrection(self.biasFilename2, applyGainCorrection=self.applyGainCorrection)
                else:
                    bias2, bro_gain2 = ccdOverScanCorrection(self.biasFilename2, applyGainCorrection=self.applyGainCorrection)

                if self.showCorrectedBias:
                    hdu = fits.PrimaryHDU(bias2)
                    hdu.writeto(corr_bias2,overwrite=True)
                    self.signal_message.emit('saved {}'.format(corr_bias2))
                self.signal_message.emit('got row-over/col-pre scan corrected bias from: {}'.format(self.biasFilename2))
            else:
                self.signal_message.emit("bias2: '{}' does not exist".format(self.biasFilename2))

        if self.flatFilename1:
            if os.path.isfile(self.flatFilename1):

                if self.useColPreScan:
                    flat1, fro_gain1 = ccdPreOverScanCorrection(self.flatFilename1, applyGainCorrection=self.applyGainCorrection)
                else:
                    flat1, fro_gain1 = ccdOverScanCorrection(self.flatFilename1, applyGainCorrection=self.applyGainCorrection)

                if self.showCorrectedFlats:
                    if bias1 is not None:
                        hdu = fits.PrimaryHDU(flat1-bias1)
                    else:
                        hdu = fits.PrimaryHDU(flat1)

                    hdu.writeto(corr_flat1,overwrite=True)
                    self.signal_message.emit('saved {}'.format(corr_flat1))
                self.signal_message.emit('got row-over/col-pre scan corrected flat from: {}'.format(self.flatFilename1))
            else:
                self.signal_message.emit("flat1: '{}' does not exist".format(self.flatFilename1))

        if self.flatFilename2:
            if os.path.isfile(self.flatFilename2):

                if self.useColPreScan:
                    flat2, fro_gain2 = ccdPreOverScanCorrection(self.flatFilename2, applyGainCorrection=self.applyGainCorrection)
                else:
                    flat2, fro_gain2 = ccdOverScanCorrection(self.flatFilename2, applyGainCorrection=self.applyGainCorrection)

                if self.showCorrectedFlats:
                    if bias2 is not None:
                        hdu = fits.PrimaryHDU(flat2-bias2)
                    else:
                        hdu = fits.PrimaryHDU(flat2)

                    hdu.writeto(corr_flat2,overwrite=True)
                    self.signal_message.emit('saved {}'.format(corr_flat2))
                self.signal_message.emit('got row-over/col-pre scan corrected flat from: {}'.format(self.flatFilename2))
            else:
                self.signal_message.emit("flat2: '{}' does not exist".format(self.flatFilename2))

        if self.applyGainCorrection:
            result_fro_gain = {
                "fro_gain1": fro_gain1,
                "fro_gain2": fro_gain2
            }
            self.signal_message.emit(f'flat row overscan derived gain will be shown.')
            self.signal_plot_fro_gain.emit(result_fro_gain)

        result = {}
        if flat1 is not None and self.showCorrectedFlats:
            result['corr_flat1'] = corr_flat1
        else:
            result['corr_flat1'] = None

        if flat2 is not None and self.showCorrectedFlats:
            result['corr_flat2'] = corr_flat2
        else:
            result['corr_flat2'] = None

        if bias1 is not None and self.showCorrectedBias:
            result['corr_bias1'] = corr_bias1
        else:
            result['corr_bias1'] = None

        if bias2 is not None and self.showCorrectedBias:
            result['corr_bias2'] = corr_bias2
        else:
            result['corr_bias2'] = None

        if flat1 is not None and flat2 is not None:
            if bias1 is not None and bias2 is not None:
                fb1 = flat1 - bias1
                fb2 = flat2 - bias2

                if self.applySelfNorm:
                    cdc = ccdDataConfig()
                    for ch in range(16):
                        idx_row1, idx_row2, idx_col1, idx_col2 = cdc.getImgChanRegion(ch=ch)
                        m1 = np.median(fb1[idx_row1:idx_row2, idx_col1:idx_col2])
                        m2 = np.median(fb2[idx_row1:idx_row2, idx_col1:idx_col2])
                        # print('m1 = {}, m2 = {}'.format(m1,m2))
                        fb1[idx_row1:idx_row2, idx_col1:idx_col2] /= m1
                        fb2[idx_row1:idx_row2, idx_col1:idx_col2] /= m2
                else:
                    m1 = np.median( sigmaclip(fb1, 3, 3 )[0] )
                    m2 = np.median( sigmaclip(fb2, 3, 3 )[0] )
                    fb1 = fb1 / m1
                    fb2 = fb2 / m2
                ratio = fb1.astype(np.float32) / fb2.astype(np.float32)
            else:
                if self.applySelfNorm:
                    cdc = ccdDataConfig()
                    for ch in range(16):
                        idx_row1, idx_row2, idx_col1, idx_col2 = cdc.getImgChanRegion(ch=ch)
                        m1 = np.median(flat1[idx_row1:idx_row2, idx_col1:idx_col2])
                        m2 = np.median(flat2[idx_row1:idx_row2, idx_col1:idx_col2])
                        # print('m1 = {}, m2 = {}'.format(m1,m2))
                        flat1[idx_row1:idx_row2, idx_col1:idx_col2] /= m1
                        flat2[idx_row1:idx_row2, idx_col1:idx_col2] /= m2
                else:
                    m1 = np.median( sigmaclip(flat1, 3, 3 )[0] )
                    m2 = np.median( sigmaclip(flat2, 3, 3 )[0] )
                    flat1 = flat1 / m1
                    flat2 = flat2 / m2

                ratio = flat1.astype(np.float32) / flat2.astype(np.float32)

            if self.correctExpTime:
                # expTime1 = float(self.ui.lineEdit_flat1ExpTime.text())
                # expTime2 = float(self.ui.lineEdit_flat2ExpTime.text())
                ratio *= (self.expTime2/self.expTime1)

            if self.anchorToZero:
                ratio = ratio - 1.0

            hdu = fits.PrimaryHDU(ratio)
            ffname = '{}/flat_ratio.fits'.format(self.wkdir)
            result['flat_ratio_ffname'] = ffname
            hdu.writeto(ffname, overwrite=True)

            result['ratio18'] = None
            result['ratio916'] = None
            if self.plotFlatRatio:
                nrow = ratio.shape[0]
                ratio18 = np.median( ratio[500:nrow//2-500, :], axis=0 )
                ratio916= np.median( ratio[nrow//2+500:nrow-500, :], axis=0 )
                result['ratio18'] = ratio18
                result['ratio916'] = ratio916

        self.signal_result.emit(result)
        self.signal_finished.emit()

    # def resume(self):
    #     self.mutex.lock()
    #     self._is_paused = False
    #     self.condition.wakeAll()
    #     self.mutex.unlock()

    # def cancel(self):
    #     self.resume()

class WorkerThread_processPTC(QThread):
    """
    This class is used to process PTC data
    """
    # 给主线程发送消息
    signal_message          = Signal(str)
    signal_result           = Signal(dict)
    signal_plot_gain        = Signal(dict)  #
    signal_finished         = Signal()

    def __init__(
            self,
            ccd_idx = '00',
            outdir = None,
            blist = None,   # bias list
            plist1 = None,   # ptc flat list
            plist2 = None,   # ptc flat list
            tlist = None,   # exp time list
            subtractBias = False,
            applyFlatCorrection = False,
            sigmaClip = 4,
            exportPTCFigs = False,
            exportPRCFigs = False,
            plotGain = False,
            useColPreScan = False,
            parent=None
        ):
        super().__init__(parent)
        self.ccd_idx            = ccd_idx
        self.outdir             = outdir
        self.blist              = blist
        self.plist1             = plist1
        self.plist2             = plist2
        self.tlist              = tlist
        self.subtractBias       = subtractBias
        self.applyFlatCorrection= applyFlatCorrection
        self.sigmaClip          = sigmaClip
        self.exportPTCFigs      = exportPTCFigs
        self.exportPRCFigs      = exportPRCFigs
        self.plotGain           = plotGain
        self.useColPreScan      = useColPreScan
        self.mutex              = QMutex()
        self.condition          = QWaitCondition()
        self.queue              = queue.Queue()

        self.data_bias = None
        self.res_flux = None
        self.res_var  = None
        self.res_gain = None

        self.expTime = None
        self.gainResult = None

    def run(self):
        self.signal_message.emit(f"start processing PTC ...")
        self.signal_message.emit(f"output dir: {self.outdir}")

        if os.path.isdir(self.outdir+str(self.ccd_idx)) is False:
            os.system(f"mkdir {self.outdir+str(self.ccd_idx)}")
            self.signal_message.emit(f"mkdir {self.outdir+str(self.ccd_idx)}")
        else:
            self.signal_message.emit(f"{self.outdir+str(self.ccd_idx)} already exists.")

        # 暂时不使用bias
        # if self.blist is not None:
        #     self.signal_message.emit(f"loading bias ...")
        #     self.data_bias = []
        #     for i in range(len(self.blist)):
        #         b = fits.getdata(self.blist[i])
        #         self.data_bias.append(b)

        self.res_flux = []
        self.res_var  = []
        self.res_gain = []
        self.expTime  = np.array(self.tlist)

        for i in range(len(self.plist1)):
            self.signal_message.emit(f"processing {self.plist1[i]} and {self.plist2[i]}")
            # pflat1 = fits.getdata(self.plist1[i]).astype(np.float32)
            # pflat2 = fits.getdata(self.plist2[i]).astype(np.float32)

            if self.useColPreScan:
                pflat1, fro_gain1 = ccdPreOverScanCorrection(self.plist1[i], applyGainCorrection=False)
                pflat2, fro_gain2 = ccdPreOverScanCorrection(self.plist2[i], applyGainCorrection=False)
            else:
                pflat1, fro_gain1 = ccdPreOverScanCorrection(self.plist1[i], applyGainCorrection=False)
                pflat2, fro_gain2 = ccdPreOverScanCorrection(self.plist2[i], applyGainCorrection=False)

            tf, tv , tg = estimateGainUsingCorrectedFlats(pflat1, pflat2,
                                    sClip=self.sigmaClip,
                                    applyFlatCorrection=self.applyFlatCorrection)

            self.res_flux.append(tf)
            self.res_var.append(tv)
            self.res_gain.append(tg)

        self.res_flux = np.array(self.res_flux).transpose()
        self.res_var  = np.array(self.res_var).transpose()
        self.res_gain = np.array(self.res_gain).transpose()

        result = {}

        result['flux'] = self.res_flux
        result['var']  = self.res_var
        result['gain'] = self.res_gain
        result['expTime'] = self.expTime

        # 将处理后的结果返回至主线程
        self.signal_result.emit(result)
        self.signal_finished.emit()
